import open3d as o3d
from utils import *
import torch
from segmentator import segment_mesh
import trimesh
import numpy as np
import time

TEST_NUM = '004'
# PCD_PATH = './reconstruction/original/01-00-01.pcd'
# pcd = o3d.io.read_point_cloud(PCD_PATH)

# # PLY_OUT_PATH = possion_ply(pcd)
# PLY_OUT_PATH = generate_ply_rec(pcd, TEST_NUM)

# Load file
PLY_OUT_PATH = './reconstruction/ply/001_ball.ply'
mesh = trimesh.load_mesh(PLY_OUT_PATH)

# Arguments
segfile = './results/json/seg/'
aggfile = './results/json/agg/'
alignfile = './reconstruction/alignment/' + TEST_NUM + '_align.txt'
scanId = TEST_NUM
kThresh = 2 
segMinVerts = 10000
original_labels = ['door', 'wall', 'floor', 'wall']

# Segment on mesh
vertices = torch.from_numpy(mesh.vertices.astype(np.float32))
faces = torch.from_numpy(mesh.faces.astype(np.int64))
ind = segment_mesh(vertices, faces, kThresh, segMinVerts)
labels, labels_unique, color_table= generate_labels(ind, original_labels)
write_to_json_agg(aggfile, scanId, labels_unique)

# Save txt and json
# np.savetxt('./results/segment/result_mesh_label_' + TEST_NUM + '_' + str(kThresh) + '_' + str(segMinVerts) + '.txt', 
#            torch.cat([vertices, color_table], dim=1).numpy())
# write_to_json_seg(segfile, scanId, kThresh, segMinVerts, ind)



# DOWN_PCD_PATH = './reconstruction/downpcd/001_normal_down.pcd'
# downpcd = down_sample_voxel(pcd, TEST_NUM)
# downpcd = o3d.io.read_point_cloud(DOWN_PCD_PATH)

# np.savetxt('./results/tensor/vert' + TEST_NUM + '.txt', vertices.numpy())
# np.savetxt('./results/tensor/ind' + TEST_NUM + '.txt', ind.numpy())
'''
labels: wall(1) floor(3) door(5) 
'''